#include <utility>
#include <assert.h>

namespace xyx
{
    template<class K, class V>
    struct AVLTreeNode
    {
        pair<K,V> _kv;
        AVLTreeNode<K,V>* _left;
        AVLTreeNode<K,V>* _right;
        AVLTreeNode<K,V>* _parent;
        int _bf;
        
        AVLTreeNode(const pair<K,V>& kv)
        :_kv(kv)
        ,_left(nullptr)
        ,_right(nullptr)
        ,_parent(nullptr)
        ,_bf(0)
        {}
    };

    template<class K, class V>
    class AVLTree
    {
        typedef AVLTreeNode<K, V> Node;
    public:
        bool Insert(const pair<K,V>& kv)
        {
            //第一次插入数据
            if(_root == nullptr)
            {
                _root = new Node(kv);
                return true;
            }
            
            //找到合适的位置插入
            Node* parent = nullptr;
            Node* cur = _root;
            
            while(cur)
            {
                if(cur->_kv.first < kv.first)
                {
                    parent = cur;
                    cur = cur->_right;
                }
                else if(cur->_kv.first > kv.first)
                {
                    parent = cur;
                    cur = cur->_left;
                }
                else
                {
                    //key已经存在
                    return false;
                }
            }
            
            //走到这里表示找到合适的位置插入
            cur = new Node(kv);
            if(parent->_kv.first < kv.first)
            {
                parent->_right = cur;
            }
            else
            {
                parent->_left = cur;
            }
            
            cur->_parent = parent;
            
            //引入平衡因子部分
            //更新平衡因子
            while(parent)
            {
                if(cur == parent->_left)
                    parent->_bf--;
                else if(cur == parent->_right)
                    parent->_bf++;
                
                //平衡因子为0，直接退出
                if(parent->_bf == 0)
                    break;
                else if(parent->_bf == 1 || parent->_bf == -1)
                {
                    //继续向上更新
                    cur = parent;
                    parent = parent->_parent;
                }
                else if(parent->_bf == 2 || parent->_bf == -2)
                {
                    //需要旋转
                    //右子树需要左旋转
                    if(parent->_bf == 2 && cur->_bf == 1)
                    {
                        RotateL(parent);
                    }
                    else if(parent->_bf == -2 && cur->_bf == -1)
                    {
                        RotateR(parent);
                    }
                    else if(parent->_bf == 2 && cur->_bf == -1)
                    {
                        RotateRL(parent);
                    }
                    else if(parent->_bf == -2 && cur->_bf == 1)
                    {
                        RotateLR(parent);
                    }
                    break;
                }
            }
            
            return true;
        }
        
        void RotateL(Node* parent)
        {
            //将curLeft连接到parent的右子树中
            //将parent连接到cur的左子树中
            //修改cur、parent、curLeft的_parent
            
            Node* cur = parent->_right;
            Node* curLeft = cur->_left;
            parent->_right = curLeft;
            
            // curLeft可能为空
            if(curLeft)
                curLeft->_parent = parent;
            
            cur->_left = parent;
            
            //提前保存parent的_parent避免节点丢失
            Node* ppnode = parent->_parent;
            parent->_parent = cur;
            

            //边界情况，修改_root
            if(parent == _root)
            {
                _root = cur;
                cur->_parent = nullptr;
            }
            else
            {
                if(ppnode->_left == parent)
                    ppnode->_left = cur;
                else
                    ppnode->_right = cur;
                
                cur->_parent = ppnode;
            }
            
            //更新平衡因子
            parent->_bf = cur->_bf = 0;
        }
        
        void RotateR(Node* parent)
        {
            //将cur的右子树连接到parent的左子数
            //让parent连接到cur的右子树
            //修改parent、cur、curRight的_parent
            
            Node* cur = parent->_left;
            Node* curRight = cur->_right;
            
            parent->_left = curRight;
            
            //curRight可能为空
            if(curRight)
                curRight->_parent = parent;
            
            //避免parent->_parent节点丢失
            Node* ppnode = parent->_parent;
            
            parent->_parent = cur;
            cur->_right = parent;
            
            //边界情况
            if(parent == _root)
            {
                _root = cur;
                cur->_parent = nullptr;
            }
            else
            {
                if(ppnode->_left == parent)
                    ppnode->_left = cur;
                else
                    ppnode->_right = cur;

                cur->_parent = ppnode;
            }
            
            
            //修改平衡因子
            cur->_bf = parent->_bf = 0;
        }
        
        void RotateRL(Node* parent)
        {
            Node* cur = parent->_right;
            Node* curLeft = cur->_left;
            int bf = curLeft->_bf;
            
            RotateR(parent->_right);
            RotateL(parent);
            
            if(bf == 0)
            {
                cur->_bf = 0;
                parent->_bf = 0;
                curLeft->_bf = 0;
            }
            else if(bf == 1)
            {
                cur->_bf = 0;
                parent->_bf = -1;
                curLeft->_bf = 0;
            }
            else if(bf == -1)
            {
                cur->_bf = 1;
                parent->_bf = 0;
                curLeft->_bf = 0;
            }
            else
            {
                assert(false);
            }
        }
        
        void RotateLR(Node* parent)
        {
            Node* cur = parent->_left;
            Node* curRight = cur->_right;
            int bf = curRight->_bf;
            
            RotateL(parent->_left);
            RotateR(parent);
            
            if(bf == 0)
            {
                cur->_bf = 0;
                curRight->_bf = 0;
                parent->_bf = 0;
            }
            else if(bf == 1)
            {
                cur->_bf = -1;
                parent->_bf = 0;
                curRight->_bf = 0;
            }
            else if(bf == -1)
            {
                cur->_bf = 0;
                parent->_bf = 1;
                curRight->_bf = 0;
            }
            else
            {
                assert(false);
            }
        }

        int Height()
        {
            return Height(_root);
        }

        int Height(Node* root)
        {
            //处理边界情况
            if (root == nullptr)
                return 0;

            //采用后序遍历的思路进行分治
            //当前节点的高度取决于左子树高度和右子树高度较大的那一个高度
            int leftHeight = Height(root->_left);
            int rightHeight = Height(root->_right);
            
            return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
        }

        bool IsBalance()
        {
            return IsBalance(_root);
        }

        bool IsBalance(Node* root)
        {
            //处理边界情况
            if (root == nullptr)
                return true;

            //求当前节点的左右子树的高度
            int leftHight = Height(root->_left);
            int rightHight = Height(root->_right);

            //判断当前节点是否满足AVL树的性质
            if (rightHight - leftHight != root->_bf)
            {
                cout << "平衡因子异常:" << root->_kv.first << "->" << root->_bf << endl;
                return false;
            }

            //当前节点的左右子树高度差的绝对值小于2 再进行子问题递归
            return abs(rightHight - leftHight) < 2
                && IsBalance(root->_left)
                && IsBalance(root->_right);
        }

    private:
        Node* _root = nullptr;
    };

}

